import os
import sys
import time
import random
import string
import logging
import threading
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from datetime import datetime
from flask import Flask, request, jsonify, make_response
from flask_cors import CORS
import requests
from dotenv import load_dotenv
import psutil

# 加载环境变量
load_dotenv()

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('server.log'),
        logging.StreamHandler()
    ]
)

app = Flask(__name__)
CORS(app)

# 配置
OLLAMA_API_URL = os.getenv('OLLAMA_API_URL', 'http://localhost:11434')
DEFAULT_PORT = int(os.getenv('DEFAULT_PORT', '4433'))
TOKEN_LENGTH = int(os.getenv('TOKEN_LENGTH', '32'))
ALLOWED_ORIGINS = os.getenv('ALLOWED_ORIGINS', '*')
RATE_LIMIT_PER_DAY = int(os.getenv('RATE_LIMIT_PER_DAY', '200'))
RATE_LIMIT_PER_HOUR = int(os.getenv('RATE_LIMIT_PER_HOUR', '50'))
HEALTH_CHECK_INTERVAL = int(os.getenv('HEALTH_CHECK_INTERVAL', '60'))
HEALTH_CHECK_TIMEOUT = int(os.getenv('HEALTH_CHECK_TIMEOUT', '5'))

# 邮件配置
SMTP_SERVER = os.getenv('SMTP_SERVER', 'smtp.gmail.com')
SMTP_PORT = int(os.getenv('SMTP_PORT', '587'))
SMTP_USERNAME = os.getenv('SMTP_USERNAME', '')
SMTP_PASSWORD = os.getenv('SMTP_PASSWORD', '')
NOTIFICATION_EMAIL = os.getenv('NOTIFICATION_EMAIL', '')

# 存储访问令牌和请求计数
tokens = {}
request_counts = {}
last_health_check = time.time()
is_healthy = True

def generate_token():
    """生成随机访问令牌"""
    return ''.join(random.choices(string.ascii_letters + string.digits, k=TOKEN_LENGTH))

def send_email(subject, body):
    """发送邮件通知"""
    if not SMTP_USERNAME or not SMTP_PASSWORD or not NOTIFICATION_EMAIL:
        logging.warning("Email configuration is incomplete. Skipping email notification.")
        return

    try:
        msg = MIMEMultipart()
        msg['From'] = SMTP_USERNAME
        msg['To'] = NOTIFICATION_EMAIL
        msg['Subject'] = subject

        msg.attach(MIMEText(body, 'plain'))

        with smtplib.SMTP(SMTP_SERVER, SMTP_PORT) as server:
            server.starttls()
            server.login(SMTP_USERNAME, SMTP_PASSWORD)
            server.send_message(msg)
        logging.info(f"Email sent successfully: {subject}")
    except Exception as e:
        logging.error(f"Failed to send email: {str(e)}")

def check_health():
    """检查服务健康状态"""
    global last_health_check, is_healthy
    while True:
        try:
            response = requests.get(f"{OLLAMA_API_URL}/api/tags", timeout=HEALTH_CHECK_TIMEOUT)
            current_status = response.status_code == 200
            if current_status != is_healthy:
                is_healthy = current_status
                status = "healthy" if is_healthy else "unhealthy"
                send_email(
                    f"Ollama Server Status Change: {status}",
                    f"The Ollama server status changed to {status} at {datetime.now()}"
                )
            last_health_check = time.time()
        except Exception as e:
            if is_healthy:
                is_healthy = False
                send_email(
                    "Ollama Server Unhealthy",
                    f"The Ollama server became unhealthy at {datetime.now()}. Error: {str(e)}"
                )
            logging.error(f"Health check failed: {str(e)}")
        time.sleep(HEALTH_CHECK_INTERVAL)

def monitor_resources():
    """监控系统资源使用情况"""
    while True:
        cpu_percent = psutil.cpu_percent()
        memory = psutil.virtual_memory()
        if cpu_percent > 90 or memory.percent > 90:
            send_email(
                "High Resource Usage Alert",
                f"High resource usage detected at {datetime.now()}:\n"
                f"CPU Usage: {cpu_percent}%\n"
                f"Memory Usage: {memory.percent}%"
            )
        time.sleep(300)  # 每5分钟检查一次

def validate_token(token):
    """验证访问令牌"""
    if token not in tokens:
        return False
    if time.time() - tokens[token]['created_at'] > 86400:  # 24小时过期
        del tokens[token]
        return False
    return True

def check_rate_limit(token):
    """检查请求频率限制"""
    if token not in request_counts:
        request_counts[token] = {'day': 0, 'hour': 0, 'last_reset': time.time()}
    
    current_time = time.time()
    counts = request_counts[token]
    
    # 重置每日计数
    if current_time - counts['last_reset'] >= 86400:
        counts['day'] = 0
        counts['hour'] = 0
        counts['last_reset'] = current_time
    
    # 重置每小时计数
    if current_time - counts['last_reset'] >= 3600:
        counts['hour'] = 0
    
    if counts['day'] >= RATE_LIMIT_PER_DAY or counts['hour'] >= RATE_LIMIT_PER_HOUR:
        return False
    
    counts['day'] += 1
    counts['hour'] += 1
    return True

@app.route('/api/generate', methods=['POST'])
def generate():
    """文本生成接口"""
    token = request.headers.get('Authorization')
    if not token or not validate_token(token):
        return jsonify({'error': 'Invalid or expired token'}), 401
    
    if not check_rate_limit(token):
        return jsonify({'error': 'Rate limit exceeded'}), 429
    
    try:
        response = requests.post(f"{OLLAMA_API_URL}/api/generate", json=request.json)
        return jsonify(response.json())
    except Exception as e:
        logging.error(f"Generate request failed: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/chat', methods=['POST'])
def chat():
    """对话接口"""
    token = request.headers.get('Authorization')
    if not token or not validate_token(token):
        return jsonify({'error': 'Invalid or expired token'}), 401
    
    if not check_rate_limit(token):
        return jsonify({'error': 'Rate limit exceeded'}), 429
    
    try:
        response = requests.post(f"{OLLAMA_API_URL}/api/chat", json=request.json)
        return jsonify(response.json())
    except Exception as e:
        logging.error(f"Chat request failed: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/embeddings', methods=['POST'])
def embeddings():
    """嵌入向量接口"""
    token = request.headers.get('Authorization')
    if not token or not validate_token(token):
        return jsonify({'error': 'Invalid or expired token'}), 401
    
    if not check_rate_limit(token):
        return jsonify({'error': 'Rate limit exceeded'}), 429
    
    try:
        response = requests.post(f"{OLLAMA_API_URL}/api/embeddings", json=request.json)
        return jsonify(response.json())
    except Exception as e:
        logging.error(f"Embeddings request failed: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/health', methods=['GET'])
def health():
    """健康检查接口"""
    return jsonify({
        'status': 'healthy' if is_healthy else 'unhealthy',
        'last_check': last_health_check,
        'ollama_status': 'healthy' if is_healthy else 'unhealthy'
    })

@app.route('/token', methods=['GET'])
def get_token():
    """获取新的访问令牌"""
    token = generate_token()
    tokens[token] = {'created_at': time.time()}
    
    # 发送邮件通知
    send_email(
        "New Ollama Linker Token Generated",
        f"A new access token has been generated at {datetime.now()}:\n\n"
        f"Token: {token}\n\n"
        f"This token will expire in 24 hours."
    )
    
    return jsonify({'token': token})

if __name__ == '__main__':
    port = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT
    
    # 启动健康检查线程
    health_thread = threading.Thread(target=check_health, daemon=True)
    health_thread.start()
    
    # 启动资源监控线程
    monitor_thread = threading.Thread(target=monitor_resources, daemon=True)
    monitor_thread.start()
    
    logging.info(f"Starting Ollama Linker server on port {port}")
    app.run(host='0.0.0.0', port=port) 